from gensim.models import Word2Vec
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# 示例：医疗领域语料
medical_sentences = [
    ['患者', '接受', 'CT', '检查'], 
    ['MRI', '用于', '诊断', '病变'],
    ['心电图', '记录', '心脏', '活动'],
    ['医生', '建议', '进行', 'MRI', '扫描'],
    ['CT', '和', 'MRI', '都是', '成像', '技术'],
    # 新增句子
    ['CT', '扫描', '可以', '检测', '肿瘤'],
    ['MRI', '成像', '提供', '更清晰', '图像'],
    ['CT', '和', 'MRI', '在', '放射科', '常用'],
    ['医生', '根据', 'CT', '和', 'MRI', '结果', '诊断'],
    ['CT', '检查', '和', 'MRI', '检查', '各有', '优势']
]

# 增加训练轮次
medical_model = Word2Vec(
    medical_sentences, 
    vector_size=100,      # 增加向量维度
    window=5,             # 增大上下文窗口
    min_count=1, 
    workers=4,
    epochs=50,            # 增加训练轮次
    negative=20,          # 增加负采样数量
    hs=1                  # 使用层次softmax
)

# 查看CT和MRI的嵌入向量
ct_vector = medical_model.wv["CT"]
mri_vector = medical_model.wv["MRI"]
print("医疗领域中 'CT' 的嵌入向量:")
print(ct_vector[:5], "...") # 展示前5个维度
print("\n医疗领域中 'MRI' 的嵌入向量:")
print(mri_vector[:5], "...")

# 计算CT和MRI的相似性
similarity_ct_mri = cosine_similarity([ct_vector], [mri_vector])[0][0]
print(f"\n'CT' 和 'MRI' 在医疗领域语料中的相似性: {similarity_ct_mri:.4f}")

# 与通用语料模型的对比
# 通用语料（假设已经训练好的Word2Vec模型）
general_sentences = [
    ['猫', '喜欢', '晒太阳'],
    ['狗', '是', '人类', '最好的', '朋友'],
    ['猫', '和', '狗', '是', '常见', '的', '宠物']
]

general_model = Word2Vec(general_sentences, vector_size=50, window=3, min_count=1, workers=4)

# 在通用模型中计算CT和MRI的相似性
try:
    general_ct_vector = general_model.wv["CT"]
    general_mri_vector = general_model.wv["MRI"]
    general_similarity_ct_mri = cosine_similarity([general_ct_vector], [general_mri_vector])[0][0]
    print(f"\n通用语料中 'CT' 和 'MRI' 的相似性: {general_similarity_ct_mri:.4f}")
except KeyError as e:
    print("\n通用语料中缺少 'CT' 或 'MRI', 无法计算相似性")

# 结果分析
# 使用医疗领域数据评估其他专业术语的相似性
similarity_mri_scan = cosine_similarity([mri_vector], [medical_model.wv["扫描"]])[0][0]
print(f"\n'MRI' 和 '扫描' 在医疗领域语料中的相似性: {similarity_mri_scan:.4f}")